{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Orbital Transfer via Trajectory Optimization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib backend (to notebook, if possible, or inline).  \n",
    "from underactuated.jupyter import setup_matplotlib_backend\n",
    "plt_is_interactive = setup_matplotlib_backend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# pydrake imports\n",
    "from pydrake.all import (Variable, SymbolicVectorSystem, DiagramBuilder,\n",
    "                         LogOutput, Simulator, ConstantVectorSource,\n",
    "                         MathematicalProgram, Solve, SnoptSolver, PiecewisePolynomial)\n",
    "\n",
    "# increase default size matplotlib figures\n",
    "from matplotlib import rcParams\n",
    "rcParams['figure.figsize'] = (8, 5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Problem Description\n",
    "\n",
    "In this problem we'll solve for trajectories that efficiently transfer a rocket from the Earth to Mars, while avoiding a huge cloud of asteroids.\n",
    "We'll make a few simplifying assumptions about the world we're in:\n",
    "- the galaxy is 2D,\n",
    "- the planets are point masses and do not move,\n",
    "- planet orbits are circular.\n",
    "\n",
    "We are not satisfied with the current trajectory we get from the nonlinear optimizer.\n",
    "Your role is to improve our trajectory by:\n",
    "- enforcing state and input limits,\n",
    "- minimizing the fuel consumption, and, more importantly,\n",
    "- avoiding asteroids!\n",
    "\n",
    "You'll be given more details on this later in the notebook.\n",
    "\n",
    "Let's start deciding the units of measurements we'll use in this problem: this is a fundamental step to prevent the solver from having numeric issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dictionary of functions to convert the units of the problem data\n",
    "# the fist argument is the numeric value we want to convert\n",
    "# the second argument is the unit power\n",
    "# e.g., m = 2 and power = 2 for square meters\n",
    "unit_converter = {\n",
    "    'mass':   lambda m, power=1 : m / 1e3 ** power, # takes kilos, returns tons\n",
    "    'length': lambda l, power=1 : l / 1e11 ** power, # takes meters, returns hundreds of gigameters\n",
    "    'time':   lambda t, power=1 : t / (60 * 60 * 24 * 30 * 12)  ** power, # takes seconds, returns years\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model of the 2D Universe\n",
    "\n",
    "Let's start by setting up the rocket.\n",
    "We're lucky enough to have in our hands a [SpaceX Falcon 9 rocket](https://www.spacex.com/falcon9).\n",
    "\n",
    "Unfortunately, we spent most of the budget buying the rocket, and we can't afford much fuel.\n",
    "For this reason we put ridiculously tight limits on the maximum thrust and velocity.\n",
    "The plan is to arrive to Mars in 50 years...\n",
    "\n",
    "(The true reason for these tiny limits is a mixture of modeling and numeric issues. If you feel like, give the true [Falcon 9 specs](https://www.spacex.com/falcon9) a try!)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# simple class to store the rocket data\n",
    "class Rocket(object):\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        mass,           # mass of the rocket in kg\n",
    "        thrust_limit,   # max norm of the thrust in kg * m * s^-2\n",
    "        velocity_limit  # max norm of the velocity in m * s^-1\n",
    "    ):\n",
    "        \n",
    "        # store mass using the scaled units\n",
    "        self.mass = unit_converter['mass'](mass)\n",
    "        \n",
    "        # store thrust limit converting the units one by one\n",
    "        thrust_units = [('mass', 1), ('length', 1), ('time', -2)]\n",
    "        for (quantity, power) in thrust_units:\n",
    "            thrust_limit = unit_converter[quantity](thrust_limit, power)\n",
    "        self.thrust_limit = thrust_limit\n",
    "        \n",
    "        # store velocity limit converting the units one by one\n",
    "        velocity_units = [('length', 1), ('time', -1)]\n",
    "        for (quantity, power) in velocity_units:\n",
    "            velocity_limit = unit_converter[quantity](velocity_limit, power)\n",
    "        self.velocity_limit = velocity_limit\n",
    "        \n",
    "# instantiate the rocket\n",
    "rocket = Rocket(\n",
    "    5.49e5, # mass of Falcon 9 in kg\n",
    "    .25,    # very small thrust limit in kg * m * s^-2\n",
    "    170,    # very small velocity limit in m * s^-1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to code the `Planet` class.\n",
    "The Earth and Mars will be instances of this class, as well as the asteroids.\n",
    "\n",
    "This class has a parameter called `orbit`.\n",
    "For the Earth, this is the radius of the orbit from which our mission starts.\n",
    "For Mars, this is the radius of the target orbit.\n",
    "For the asteroids, this is the radius of a \"danger area\" centered at the asteroid itself that we want our rocket to avoid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# each planet/asteroid in the problem must be an instance of this class\n",
    "class Planet(object):\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        name,          # string with the name of the planet\n",
    "        color,         # color of the planet for plots\n",
    "        mass,          # mass of the planet in kg\n",
    "        position,      # position of the planet in the 2d universe in m\n",
    "        orbit,         # radius of the orbit in m\n",
    "        radius=np.nan, # radius of the planet in m (optional)\n",
    "    ):\n",
    "        \n",
    "        # store the data using the scaled units\n",
    "        self.name = name\n",
    "        self.mass = unit_converter['mass'](mass)\n",
    "        self.position = unit_converter['length'](position)\n",
    "        self.radius = unit_converter['length'](radius)\n",
    "        self.orbit = unit_converter['length'](orbit)\n",
    "        self.color = color\n",
    "        \n",
    "# planet Earth: https://en.wikipedia.org/wiki/Earth\n",
    "earth = Planet(\n",
    "    'Earth',                # name of the planet\n",
    "    'green',                # color for plot\n",
    "    5.972e24,               # mass in kg\n",
    "    np.array([2.25e11, 0]), # (average) distance wrt Mars in m\n",
    "    2e10,                   # orbit radius in m (chosen \"big enough\" for the plots)\n",
    "    6.378e6,                # planet radius in m\n",
    ")\n",
    "\n",
    "# planet Mars: https://en.wikipedia.org/wiki/Mars\n",
    "mars = Planet(\n",
    "    'Mars',      # name of the planet\n",
    "    'red',       # color for plot\n",
    "    6.417e23,    # mass in kg\n",
    "    np.zeros(2), # Mars is chosen as the origin of our 2D universe\n",
    "    1.5e10,      # orbit radius in m\n",
    "    3.389e6,     # radius in m\n",
    ")\n",
    "\n",
    "# asteroids with random data in random positions\n",
    "np.random.seed(0)\n",
    "n_asteroids = 10\n",
    "asteroids = []\n",
    "for i in range(n_asteroids):\n",
    "    mass = np.abs(np.random.randn()) * 5e22\n",
    "    orbit = mass / 5e12\n",
    "    earth_from_mars = unit_converter['length'](earth.position, -1)\n",
    "    asteroid_from_mars = np.random.randn(2) * 3e10 + earth_from_mars / 2\n",
    "    asteroids.append(\n",
    "        Planet(\n",
    "            f'Asteroid_{i}',    # name of the planet\n",
    "            'brown',            # color for plot\n",
    "            mass,               # mass in kg\n",
    "            asteroid_from_mars, # distance from Mars in m\n",
    "            mass / 5e12,        # radius danger area in m\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's put all the pieces together in ours 2D universe.\n",
    "The class `Universe` also implements some utility functions that will be helpful when setting up the optimization problem.\n",
    "For example, for the piece of code you'll be asked to write, you might consider using the `position_wrt_planet` method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# main class of the notebook\n",
    "# it collects the rocket, the planets, and all the asteroids\n",
    "# implements utility functions needed to write the trajopt\n",
    "class Universe(object):\n",
    "    \n",
    "    def __init__(\n",
    "        self,\n",
    "        rocket, # instance of Rocket\n",
    "        planets # list of instances of Planet\n",
    "    ):\n",
    "        \n",
    "        # store data\n",
    "        self.rocket = rocket\n",
    "        self.planets = planets\n",
    "        \n",
    "        # gravitational constant in m^3 * kg^-1 * s^-2\n",
    "        self.G = 6.67e-11\n",
    "        \n",
    "        # gravitational constant in the scaled units\n",
    "        G_units = [('length', 3), ('mass', -1), ('time', -2)]\n",
    "        for (quantity, power) in G_units:\n",
    "            self.G = unit_converter[quantity](self.G, power)\n",
    "            \n",
    "    # given the planet name, returns the Planet instance\n",
    "    def get_planet(self, name):\n",
    "        \n",
    "        # loop through the planets in the universe\n",
    "        for planet in self.planets:\n",
    "            if planet.name == name:\n",
    "                return planet\n",
    "            \n",
    "        # in case no planet has the given name\n",
    "        print(name + ' is not in the Universe!')\n",
    "        \n",
    "    # computes  2D distance vector between the rocket and a planet,\n",
    "    # given the rocket state and the planet name\n",
    "    def position_wrt_planet(self, state, name):\n",
    "        \n",
    "        # rocket position wrt to the planet position\n",
    "        planet = self.get_planet(name)\n",
    "        p = state[:2] - planet.position\n",
    "\n",
    "        return p\n",
    "            \n",
    "    # computes the rocket acceleration due to a planet\n",
    "    def acceleration_from_planet(self, state, name):\n",
    "\n",
    "        # distance from the planet\n",
    "        p = self.position_wrt_planet(state, name)\n",
    "        d = p.dot(p) ** .5\n",
    "        \n",
    "        # 2d acceleration vector\n",
    "        planet = self.get_planet(name)\n",
    "        a = - self.G * planet.mass / d ** 3  * p\n",
    "\n",
    "        return a\n",
    "    \n",
    "    # right-hand side of the rocket continuous-time dynamics\n",
    "    # in the form state_dot = f(state, thrust)\n",
    "    # (thrust is a 2D vector with the horizontal and vertical thrusts)\n",
    "    def rocket_continuous_dynamics(self, state, thrust):\n",
    "\n",
    "        # thrust acceleration\n",
    "        a = thrust / self.rocket.mass\n",
    "\n",
    "        # accelerations due to the planets\n",
    "        for planet in self.planets:\n",
    "            a = a + self.acceleration_from_planet(state, planet.name)\n",
    "            \n",
    "        # concatenate velocity and acceleration\n",
    "        state_dot = np.concatenate((state[2:], a))\n",
    "            \n",
    "        return state_dot\n",
    "    \n",
    "    # residuals of the rocket discrete-time dynamics\n",
    "    # if the vector of residuals is zero, then this method's\n",
    "    # arguments verify the discrete-time dynamics\n",
    "    # (implements the implicit Euler integration scheme:\n",
    "    # https://en.wikipedia.org/wiki/Backward_Euler_method)\n",
    "    def rocket_discrete_dynamics(self, state, state_next, thrust, time_step):\n",
    "        \n",
    "        # continuous-time dynamics evaluated at the next time step\n",
    "        state_dot = self.rocket_continuous_dynamics(state_next, thrust)\n",
    "        \n",
    "        # implicit-Euler state update\n",
    "        residuals = state_next - state - time_step * state_dot\n",
    "        \n",
    "        return residuals\n",
    "    \n",
    "    # helper function for the trajopt problem\n",
    "    # if the vector of residuals is zero, then the state of\n",
    "    # the rocket belongs to the desired orbit of the given planet\n",
    "    # (i.e.: the rocket is on the given orbit, with zero radial\n",
    "    # velocity, and zero radial acceleration)\n",
    "    def constraint_state_to_orbit(self, state, planet_name):\n",
    "        \n",
    "        # unpack state, rocket position in relative coordinates\n",
    "        planet = self.get_planet(planet_name)\n",
    "        p = state[:2] - planet.position\n",
    "        v = state[2:]\n",
    "\n",
    "        # constraint on radial distance\n",
    "        # sets x^2 + y^2 to the orbit radius squared\n",
    "        residual_p = p.dot(p) - planet.orbit ** 2\n",
    "\n",
    "        # radial velocity must be zero\n",
    "        # sets the time derivative of x^2 + y^2 to zero\n",
    "        residual_v = p.dot(v)\n",
    "\n",
    "        # radial acceleration must be zero with zero input\n",
    "        # sets the second time derivative of x^2 + y^2 to zero\n",
    "        # why this extra constraint?\n",
    "        # knowing that the radial velocity is zero is not enough\n",
    "        # the tangential velocity must be such that the gravitational\n",
    "        # force is balanced by the centrifugal force\n",
    "        a = self.acceleration_from_planet(state, planet_name)\n",
    "        residual_a = p.dot(a) + v.dot(v)\n",
    "        \n",
    "        # gather constraint residuals\n",
    "        residuals = np.array([residual_p, residual_v, residual_a])\n",
    "\n",
    "        return residuals\n",
    "    \n",
    "    # bonus method! (not actually needed in the trajopt...)\n",
    "    # computes the gravity acceleration on the surface of a planet\n",
    "    def gravity_on_planet_surface(self, name):\n",
    "        \n",
    "        # retrieve planet\n",
    "        planet = self.get_planet(name)\n",
    "        if planet is not None:\n",
    "            \n",
    "            # if planet radius is not available\n",
    "            if np.isnan(planet.radius):\n",
    "                print(name + ' has unknown radius.')\n",
    "                return\n",
    "            \n",
    "            # use Newton's law of universal gravitation\n",
    "            g = self.G * planet.mass / planet.radius ** 2\n",
    "            \n",
    "            # use the converter the other way around\n",
    "            # to express g in MKS\n",
    "            g_inverse_units = [('length', -1), ('time', 2)]\n",
    "            for (quantity, power) in g_inverse_units:\n",
    "                g = unit_converter[quantity](g, power)\n",
    "                \n",
    "            # print the result\n",
    "            print('Gravity acceleration on ' + name + f' is {g} m/s^2.')\n",
    "            \n",
    "            \n",
    "# instantiate universe\n",
    "planets = [earth, mars] + asteroids\n",
    "universe = Universe(rocket, planets)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Double Check the Physics\n",
    "\n",
    "Just to be sure that things make sense, let's compute the gravity accelerations on the surface of the planets in our 2D universe.\n",
    "Are these values close to what you expect?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "universe.gravity_on_planet_surface('Earth')\n",
    "universe.gravity_on_planet_surface('Mars')\n",
    "universe.gravity_on_planet_surface('Jupiter')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot Functions\n",
    "\n",
    "Here are some plot functions that will help us visualizing the results.\n",
    "You'll find these useful when improving the trajectory as explained below.\n",
    "As always, no need to fully understand what's going on here, unless you're a `matplotlib` enthusiast..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# helper function that plots a circle centered at\n",
    "# the given point and with the given radius\n",
    "def plot_circle(center, radius, *args, **kwargs):\n",
    "    \n",
    "    # discretize angle\n",
    "    angle = np.linspace(0, 2*np.pi)\n",
    "    \n",
    "    # plot circle\n",
    "    plt.plot(\n",
    "        center[0] + radius * np.cos(angle),\n",
    "        center[1] + radius * np.sin(angle),\n",
    "        *args,\n",
    "        **kwargs\n",
    "    )\n",
    "\n",
    "# function that draws the state-space trajectory of the rocket\n",
    "# including the planets and the asteroids\n",
    "def plot_state_trajectory(trajectory, universe):\n",
    "    \n",
    "    for planet in universe.planets:\n",
    "        \n",
    "        # plot planets\n",
    "        plt.scatter(*planet.position, s=100, c=planet.color)\n",
    "        plt.text(*planet.position, planet.name)\n",
    "    \n",
    "        # plot orbits\n",
    "        if not np.isnan(planet.orbit):\n",
    "            if planet.name == 'Asteroid_1':\n",
    "                orbit_label = 'Asteroid danger area'\n",
    "            elif planet.name[:8] == 'Asteroid':\n",
    "                orbit_label = None\n",
    "            else:\n",
    "                orbit_label = planet.name + ' orbit'\n",
    "            plot_circle(\n",
    "                planet.position,\n",
    "                planet.orbit,\n",
    "                label=orbit_label,\n",
    "                color=planet.color,\n",
    "                linestyle='--'\n",
    "            )\n",
    "    \n",
    "    # plot rocket trajectory\n",
    "    plt.plot(trajectory.T[0], trajectory.T[1], color='k', label='Rocket trajectory')\n",
    "    plt.scatter(trajectory[0,0], trajectory[0,1], color='k')\n",
    "    \n",
    "    # misc settings\n",
    "    length_unit = unit_converter['length'](1)\n",
    "    plt.xlabel('{:.0e} meters'.format(length_unit))\n",
    "    plt.ylabel('{:.0e} meters'.format(length_unit))\n",
    "    plt.grid(True)\n",
    "    plt.gca().set_aspect('equal')\n",
    "    \n",
    "    # legend\n",
    "    n_legend = len(plt.gca().get_legend_handles_labels()[0])\n",
    "    plt.legend(\n",
    "        loc='upper center',\n",
    "        ncol=int(n_legend / 2),\n",
    "        bbox_to_anchor=(.5, 1.25),\n",
    "        fancybox=True,\n",
    "        shadow=True\n",
    "    )\n",
    "\n",
    "# function that plots the norm of the rocket thrust and\n",
    "# velocity normalized on their maximum value\n",
    "def plot_rocket_limits(rocket, thrust, state):\n",
    "    \n",
    "    # reconstruct time vector\n",
    "    time_steps = thrust.shape[0]\n",
    "    time = np.linspace(0, time_steps, time_steps + 1)\n",
    "    \n",
    "    # plot maximum norm limit\n",
    "    plt.plot(time, np.ones(time_steps + 1), 'r--', label='Limit')\n",
    "    \n",
    "    # plot normalized thrust\n",
    "    thrust_norm = [np.linalg.norm(t) / rocket.thrust_limit for t in thrust]\n",
    "    plt.step(time, [thrust_norm[0]] + thrust_norm, label='Thrust / thrust limit')\n",
    "    \n",
    "    # plot normalized velocity\n",
    "    velocity_norm = [np.linalg.norm(v) / rocket.velocity_limit for v in state[:,2:]]\n",
    "    plt.plot(time, velocity_norm, label='Velocity / velocity limit')\n",
    "    \n",
    "    # plot limits\n",
    "    plt.xlim(0, time_steps)\n",
    "    ymax = max(1, max(thrust_norm), max(velocity_norm)) * 1.05\n",
    "    plt.ylim(0, ymax)\n",
    "    \n",
    "    # misc settings\n",
    "    plt.xlabel('Time step')\n",
    "    plt.grid(True)\n",
    "    plt.legend()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Trajectory Optimization\n",
    "\n",
    "We start writing a simple function to generate an initial guess for the trajectory of the rocket state.\n",
    "\n",
    "We just draw a straight line from the Earth to Mars.\n",
    "This does not verify our constraints, but we let the solver think about the details...\n",
    "We just need to give it a rough idea of how we want the final trajectory to look like.\n",
    "\n",
    "Note that we also add a tiny amount of randomness to the initial guess.\n",
    "This is needed since a straight line from the Earth to Mars would have the initial and the final knot points in the centers of the planets.\n",
    "There, the Gravitational force is infinite, and the solver would crash at the first iteration!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# function that interpolates two given positions of the rocket\n",
    "# velocity is set to zero for all the times\n",
    "def interpolate_rocket_state(p_initial, p_final, time_steps):\n",
    "    np.random.seed(0)\n",
    "    \n",
    "    # initial and final time and state\n",
    "    time_limits = [0., time_steps * time_interval]\n",
    "    position_limits = np.column_stack((p_initial, p_final))\n",
    "    state_limits = np.vstack((position_limits, np.zeros((2, 2))))\n",
    "    \n",
    "    # linear interpolation in state\n",
    "    state = PiecewisePolynomial.FirstOrderHold(time_limits, state_limits)\n",
    "\n",
    "    # sample state on the time grid and add small random noise\n",
    "    state_guess = np.vstack([state.value(t * time_interval).T for t in range(time_steps + 1)])\n",
    "    state_guess += np.random.rand(*state_guess.shape) * 5e-6\n",
    "    \n",
    "    return state_guess"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's finally time to write the trajectory optimization problem.\n",
    "Here is the skeleton of it.\n",
    "So far we just:\n",
    "- defined the optimization variables,\n",
    "- enforced the initial and terminal constraints,\n",
    "- required the rocket to satisfy the discretized dynamics we defined above,\n",
    "- set the state initial guess.\n",
    "\n",
    "**Note:**\n",
    "In this notebook we write the trajectory optimization as a generic `MathematicalProgram`.\n",
    "This because we want your hands to get dirty at least once, and see in detail how these optimization problems are constructed.\n",
    "Having said this, Drake provides faster and cleaner ways to write down trajectory optimization problems.\n",
    "If you are interested, have a look at the class [`DirectCollocation`](https://drake.mit.edu/doxygen_cxx/classdrake_1_1systems_1_1trajectory__optimization_1_1_direct_collocation.html) and at [this notebook](https://github.com/RussTedrake/underactuated/blob/master/examples/trajectory_optimization.ipynb).\n",
    "\n",
    "**Troubleshooting:**\n",
    "We've split the setup and the solution of the optimization problem in multiple cells.\n",
    "This is needed for you to fill in the missing pieces.\n",
    "Note that after you've done your modifications to the mathematical program, you'll need to rerun all the cells starting from the one below, where the `MathematicalProgram` is initialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# numeric parameters\n",
    "time_interval = .5 # in years\n",
    "time_steps = 100\n",
    "\n",
    "# initialize optimization\n",
    "prog = MathematicalProgram()\n",
    "\n",
    "# optimization variables\n",
    "state = prog.NewContinuousVariables(time_steps + 1, 4, 'state')\n",
    "thrust = prog.NewContinuousVariables(time_steps, 2, 'thrust')\n",
    "\n",
    "# initial orbit constraints\n",
    "for residual in universe.constraint_state_to_orbit(state[0], 'Earth'):\n",
    "    prog.AddConstraint(residual == 0)\n",
    "\n",
    "# terminal orbit constraints\n",
    "for residual in universe.constraint_state_to_orbit(state[-1], 'Mars'):\n",
    "    prog.AddConstraint(residual == 0)\n",
    "    \n",
    "# discretized dynamics\n",
    "for t in range(time_steps):\n",
    "    residuals = universe.rocket_discrete_dynamics(state[t], state[t+1], thrust[t], time_interval)\n",
    "    for residual in residuals:\n",
    "        prog.AddConstraint(residual == 0)\n",
    "    \n",
    "# initial guess\n",
    "state_guess = interpolate_rocket_state(\n",
    "    universe.get_planet('Earth').position,\n",
    "    universe.get_planet('Mars').position,\n",
    "    time_steps\n",
    ")\n",
    "prog.SetInitialGuess(state, state_guess)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is an empty cell for you to work in: detailed instructions are given below.\n",
    "For the moment just run it as it is and continue to the plots.\n",
    "\n",
    "**Troubleshooting:**\n",
    "If you run the following cell multiple times, you keep adding more and more constraints to the optimization problem.\n",
    "Hence, after you modify the code in the cell below, rerun the cell above to start with a fresh `MathematicalProgram`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# thrust limits, for all t:\n",
    "# two norm of the rocket thrust\n",
    "# lower or equal to the rocket thrust_limit\n",
    "\n",
    "# modify here\n",
    "\n",
    "# velocity limits, for all t:\n",
    "# two norm of the rocket velocity\n",
    "# lower or equal to the rocket velocity_limit\n",
    "\n",
    "# modify here\n",
    "\n",
    "# avoid collision with asteroids, for all t, for all asteroids:\n",
    "# two norm of the rocket distance from the asteroid\n",
    "# greater or equal to the asteroid orbit\n",
    "\n",
    "# modify here\n",
    "\n",
    "# minimize fuel consumption, for all t:\n",
    "# add to the objective the two norm squared of the thrust\n",
    "# multiplied by the time_interval so that the optimal cost\n",
    "# approximates the time integral of the thrust squared\n",
    "\n",
    "# modify here"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have written the optimization problem, we can solve it using the nonlinear optimization solver `Snopt`.\n",
    "\n",
    "**Troubleshooting:**\n",
    "`Snopt` is a commercial solver and requires a licence to be used.\n",
    "`Snopt` is included in the precompiled binaries of Drake.\n",
    "Therefore,  if you installed Drake using the precompiled binaries (or you are working in Colab or Binder), the following cell will run fine.\n",
    "But, if you built Drake from source, the following cell won't work unless you built Drake with the `Snopt` option enabled ([see here](https://drake.mit.edu/python_bindings.html)) and you have a `Snopt` licence on your machine.\n",
    "In that case, you can consider the alternative free solver `IPOPT` which can be imported running `from pydrake.all import IpoptSolver`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# solve mathematical program\n",
    "solver = SnoptSolver()\n",
    "result = solver.Solve(prog)\n",
    "\n",
    "# be sure that the solution is optimal\n",
    "assert result.is_success()\n",
    "\n",
    "# retrieve optimal solution\n",
    "thrust_opt = result.GetSolution(thrust)\n",
    "state_opt = result.GetSolution(state)\n",
    "\n",
    "# compute fuel consumption for the optimal trajectory\n",
    "def fuel_consumption(thrust, time_interval):\n",
    "    return time_interval * sum(t.dot(t) for t in thrust)\n",
    "print(f'Is fuel consumption {fuel_consumption(thrust_opt, time_interval)} lower than 250?')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the Optimization Results\n",
    "\n",
    "Here is a visualization of the optimal state trajectory.\n",
    "Note that we would like the trajectory of the rocket not to intersect any of the danger areas around the asteroids."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plot_state_trajectory(state_opt, universe)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are the ratios between the 2 norm of the rocket thrust and velocity and their maximum values.\n",
    "Note that we would like these ratios not to exceed one (labeled as \"Limit\" in the plot)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plot_rocket_limits(rocket, thrust_opt, state_opt)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Your Time to Code\n",
    "\n",
    "As you can see, the current trajectory has several issues.\n",
    "First of all it, collides with the asteroids.\n",
    "Second, it exceeds both the thrust and velocity limits.\n",
    "\n",
    "We left a cell with comments only in the center of the definition of the optimization problem.\n",
    "Your goal is to work on that cell to improve the rocket trajectory.\n",
    "\n",
    "You're required to add three groups of constraints:\n",
    "- **Thrust limits.**\n",
    "The class `Rocket` has the attribute `thrust_limit`.\n",
    "For each time step `t`, add a constraint that requires the norm of the rocket thrust vector to be lower or equal to the `thrust_limit`.\n",
    "- **Velocity limits.**\n",
    "For each time step `t`, add a constraint that requires the norm of the rocket velocity (third and fourth components of the state) to be lower or equal to the `velocity_limit` from the `Rocket` class.\n",
    "- **Avoid collision with asteroids.**\n",
    "Each asteroid in the `universe` has the attribute `orbit`.\n",
    "This is the radius of a danger orbit where we don't want our rocket to end up.\n",
    "For each time step `t` and for each asteroid, add a constraint that requires distance between the rocket and the asteroid to be greater or equal to the asteroid's `orbit`.\n",
    "- **Minimize fuel consumption.**\n",
    "To ensure that all these extra constraints do not lead to crazy trajectories, add an objective function that minimizes the (discrete-time approximation) of the integral of the thrust squared: $$h \\sum_{t=0}^{T-1} \\| \\mathbf{u}[t] \\|^2,$$ where $h$ is the time interval (`time_interval` in the code), $t$ is the time step, $T$ is the number of time steps (`time_steps` in the code), $\\mathbf{u}[t]$ is the 2D thrust at discrete time $t$ (`thrust[t]` in the code).\n",
    "The current fuel consumption can be verifies using the function `fuel_consumption` defined above.\n",
    "The goal for you is to keep the consumption lower than 250 (expressed in the scaled units).\n",
    "\n",
    "\n",
    "If you're struggling trying to find the right functions to complete your task, here are two great resources:\n",
    "- [Drake's MathematicalProgram documentation](https://drake.mit.edu/doxygen_cxx/classdrake_1_1solvers_1_1_mathematical_program.html).\n",
    "This is the `C++` documentation, but it's very useful to see which methods `MathematicalProgram` implements and what you need to pass to them.\n",
    "- [Drake's nonlinear-programming tutorial](https://mybinder.org/v2/gh/RobotLocomotion/drake/nightly-release?filepath=tutorials).\n",
    "Here you'll find usage examples of all the `MathematicalProgram` methods you need to complete this exercise.\n",
    "\n",
    "**Troubleshooting:**\n",
    "Nonlinear optimization solvers are extremely sensitive.\n",
    "Different OSs can give different results.\n",
    "Even worse, sometimes, enforcing the same constraints in a different order, you can make these solvers find a solution to a problem that they previously said was infeasible!\n",
    "In our Colab testings, we implemented the constraints in the order given above, and we got full score.\n",
    "However, in case you believe that your implementation is correct but you still don't get full score from the autograder, feel free to play with the initial-guess parameters (e.g. the random seed)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "You can check your work by running the following cell."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.trajopt.orbital_transfer.test_orbital_transfer import TestOrbitalTransfer\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestOrbitalTransfer], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}